# optimal_knn.py

import pickle
import matplotlib.pyplot as plt
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score
from tqdm import tqdm  # 打印进度条

# 加载手写数字数据集
digits = load_digits()
X, y = digits.data, digits.target

# 划分训练集与测试集
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# 初始化变量
best_accuracy = 0.0
best_k = 1
best_model = None
accuracies = []

# 遍历 K 值并打印进度条
for k in tqdm(range(1, 41), desc="正在测试不同的K值"):
    model = KNeighborsClassifier(n_neighbors=k)
    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)
    acc = accuracy_score(y_test, y_pred)
    accuracies.append(acc)

    if acc > best_accuracy:
        best_accuracy = acc
        best_k = k
        best_model = model

# 保存最佳模型
with open("best_knn_model.pkl", "wb") as f:
    pickle.dump(best_model, f)

print(f"\n✅ 最佳K值: {best_k}, 对应准确率: {best_accuracy:.4f}")

# 绘图
plt.figure(figsize=(10, 6))
plt.plot(range(1, 41), accuracies, marker='o', linestyle='-', label='Accuracy')
plt.axvline(x=best_k, color='red', linestyle='--', label=f'Best K = {best_k}')
plt.scatter(best_k, best_accuracy, color='red', zorder=5)
plt.text(best_k + 0.5, best_accuracy, f"K={best_k}\nAcc={best_accuracy:.4f}", color='red')
plt.title("KNN Accuracy vs K Value")
plt.xlabel("K value")
plt.ylabel("Accuracy")
plt.legend()
plt.grid(True)
plt.tight_layout()

# 保存为 PDF 文件
plt.savefig("knn_results.pdf")
plt.close()

print("📊 已保存准确率折线图为 knn_results.pdf")
print("💾 已保存最佳模型为 best_knn_model.pkl")
